"""
rag_system.py定义了RAGSystem类，整合系统的各个模块，完成从查询输入到答案生成的完整流程。
它通过查询分类选择处理路径，利用检索策略优化文档检索，并结合上下文生成最终答案。
功能概述：
    1：分类决策，通过机器学习模型区分通用知识还是专业咨询的问题
    2：上下文构建，动态检索并整合相关文档，构建上下文信息
    3：智能查询处理，根据查询内容自动选择检索策略（如HyDE、子查询、回溯、直接查询等）
    4：答案生成，检索向量库或者大模型返回给用户内容并进行润色

实现思路：
    1：初始化配置，集成向量数据库，语言模型，查询分类器和策略选择器
    2：多策略检索，实现四种检索增强技术的具体逻辑
    3：查询分类，预判问题类型决定处理流程（如通用知识查询或专业咨询）
    4：上下文构建，文档去重，筛选和格式化
    5：答案生成，利用上下文和查询结果生成最终答案
"""
import time

# 导包
from base.config import get_config
from base.logger import get_logger
from rag_qa.core.query_classifier import QueryClassifier # 导入查询分类器
from rag_qa.core.vector_store import VectorStore # 导入向量数据库
from rag_qa.core.strategy_selector import StrategySelector # 导入策略选择器
from rag_qa.core.prompts import RAGPrompts # 导入提示模板

class RAGSystem:
    def __init__(self, vector_store: VectorStore, llm):
        """
        初始化RAG系统，加载配置文件和日志记录器。
        """
        self.config = get_config()
        self.logger = get_logger()

        # 初始化向量数据库
        self.vector_store = VectorStore()

        # 初始化语言模型
        self.llm = llm

        # 初始化提示词模板
        self.prompts = RAGPrompts.rag_prompt()

        # 初始化查询分类器
        self.query_classifier = QueryClassifier()

        # 初始化策略选择器
        self.strategy_selector = StrategySelector()

        self.llmStream = llmStream

    """
    需求：实现假设问题嵌入（HyDE）检索策略
    实现思路：
    1：使用专用提示词模板生成假设答案
    2：基于假设答案进行向量搜索
    3：应用混合检索（语义+关键词）和重排序
    """
    def _retrieve_with_hyde(self, query: str, source_filter: str = None):
        self.logger.info(f"使用假设问题嵌入（HyDE）检索策略，查询：{query}")
        # 获取假设问题生成的Prompt模板
        hyde_prompt = RAGPrompts.hyde_prompt().format(query=query)

        try:
            # 调用大模型生成假设答案，并返回检索结果
            hypo_answer = self.llm(hyde_prompt).strip()
            self.logger.info(f"假设问题嵌入（HyDE）检索策略生成假设答案：{hypo_answer}")
            # 使用假设答案进行检索，并返回检索结果（检索vector）
            return self.vector_store.hybrid_search_with_rerank(
                query=hypo_answer,
                k=self.config.RETRIEVAL_K,
                source_filter=source_filter
            )
        except Exception as e:
            self.logger.error(f"假设问题嵌入（HyDE）检索策略生成假设答案时出错：{e}")
            return []

    """
    需求：动态选择检索策略并整合结果
    实现思路：
    1：未指定策略时通过策略选择器决策
    2：根据策略选择器路由到指定检索策略方法
    3：限制最终上下文文档数量，避免超出模型最大上下文长度
    """
    def retrieve_and_merge(self, query: str, source_filter: str = None, strategy: str = None):
        if not strategy:
            stratey = self.strategy_selector.select_strategy(query)
            self.logger.info(f"未指定策略,根据查询{query}选择的策略是:{stratey}")

        if stratey =="假设问题检索":
            ranked_sub_chunks = self._retrieve_with_hyde(query, source_filter)
        elif strategy =="子查询检索":
            ranked_sub_chunks = self._retrieve_with_subqueries(query,source_filter)
        elif strategy=="回溯问题检索":
            pass
        else:
            self.logger.info(f"使用直接查询策略，查询:{query}")
            ranked_sub_chunks = self.vector_store.hybrid_search_with_rerank(
                query=query,
                k=self.config.RETRIEVAL_K,
                source_filter=source_filter
            )
        self.logger.info(f"策略{strategy},检索到{len(ranked_sub_chunks)}个候选文档")
        final_context_docs = ranked_sub_chunks[:self.config.CANDIDATE_M]
        self.logger.info(f"最终选择{len(final_context_docs)}个文档作为上下文")

    """
    需求：端对端处理用户查询并生成答案
    实现思路：
    1：使用用户意图识别模型判断问题类型（通用知识/专业咨询）
    2：通用知识：直接调用LLM生成答案
    3: 专业咨询：
        3.1：根据查询内容自动选择检索策略（如HyDE、子查询、回溯、直接查询等）
        3.2：利用检索策略优化文档检索，整合相关文档
        3.3：构建上下文信息，包含查询、文档和历史消息
        3.4：调用LLM生成专业咨询答案
    """
    def generate_answer(self, query: str, source_filter: str = None):
        # 记录查询开始时间（为了后续更好的优化，将每次查询的时间记录下来）
        start_time = time.time()
        self.logger.info(f"开始处理查询：{query}，源筛选器：{source_filter}")

        # 1：使用用户意图识别模型判断问题类型（通用知识/专业咨询）
        query_type = self.query_classifier.predict_query(query)
        self.logger.info(f"查询分类结果：{query_type}，查询内容：{query}")

        # 2：通用知识：直接调用LLM生成答案
        if query_type == "通用知识":
            self.logger.info("查询类型为通用知识，直接调用LLM生成答案")
            # 格式化提示词
            prompt_input = self.prompts.format(
                context="", question=query, phone=self.config.CUSTOMER_SERVICE_PHONE
            )
            try:
                # 调用大模型生成答案
                answer = self.llm(prompt_input)
                self.logger.info(f"直接调用LLM生成答案：{answer}")
            except Exception as e:
                self.logger.error(f"调用大模型生成答案出错：{e}")
                answer = f"信息不足，无法回答，请联系人工客服，电话：{self.config.CUSTOMER_SERVICE_PHONE}"
            # 获取查询所需要的时间
            process_time = time.time() - start_time
            self.logger.info(f"t通用知识查询处理完成（耗时：{process_time:.2f}）秒，查询：{query}")
            return answer

        # 3: 专业咨询：
        self.logger.info("查询类型为专业咨询，执行RAG处理流程")
        # 3.1：根据查询内容自动选择检索策略（如HyDE、子查询、回溯、直接查询等）
        strategy = self.strategy_selector.select_strategy(query)
        self.logger.info(f"选择检索策略：{strategy}，查询内容：{query}")

        # 检索相关文档
        context_docs = self.retrieve_and_merge(query, source_filter=source_filter, strategy=strategy)
        print(context_docs)

if __name__ == '__main__':
    # 初始化向量数据库
    vector_store = VectorStore()
    # 初始化语言模型
    llm = StrategySelector().call_dashscope
    #初始化RagSystem对象
    rag_system = RAGSystem(vector_store, llm)
    answer = rag_system.generate_answer(query="AI学科的课程内容是什么?", source_filter="ai")
    # answer = rag_system.generate_answer(query="5*9等于多少?", source_filter="AI")
    print(answer)
